package com.nbsaas.boot.queue.scheduler.impl;

import com.nbsaas.boot.queue.monitor.TaskMonitor;
import com.nbsaas.boot.queue.queue.TaskQueue;
import com.nbsaas.boot.queue.scheduler.TaskScheduler;
import com.nbsaas.boot.queue.storage.TaskStorage;
import com.nbsaas.boot.queue.task.Task;
import com.nbsaas.boot.queue.task.TaskStatus;

import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

public class DefaultTaskScheduler implements TaskScheduler {
    private final TaskQueue taskQueue;
    private final TaskStorage taskStorage;
    private final TaskMonitor taskMonitor;
    private final ExecutorService executorService;
    private final ScheduledExecutorService scheduledExecutorService;
    private volatile boolean running;

    public DefaultTaskScheduler(int threadPoolSize, TaskStorage taskStorage) {
        this.taskQueue = new TaskQueue();
        this.taskStorage = taskStorage;
        this.taskMonitor = new TaskMonitor(taskStorage);
        this.executorService = Executors.newFixedThreadPool(threadPoolSize);
        this.scheduledExecutorService = Executors.newScheduledThreadPool(1);
        this.running = true;

        // 从存储中恢复待执行的任务
        restorePendingTasks();
    }

    private void restorePendingTasks() {
        taskStorage.getPendingTasks().forEach(task -> {
            taskQueue.addTask(task);
            taskMonitor.taskCreated(task);
            scheduleTaskExecution(task);
        });
    }

    @Override
    public void start() {
        // 不需要额外的启动线程，因为ScheduledExecutorService会自动处理
    }

    @Override
    public void shutdown() {
        running = false;
        executorService.shutdown();
        scheduledExecutorService.shutdown();
        try {
            if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
            if (!scheduledExecutorService.awaitTermination(60, TimeUnit.SECONDS)) {
                scheduledExecutorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            scheduledExecutorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    @Override
    public void scheduleTask(Task task) {
        taskStorage.saveTask(task);
        taskQueue.addTask(task);
        taskMonitor.taskCreated(task);
        scheduleTaskExecution(task);
    }

    private void scheduleTaskExecution(Task task) {
        long delay = java.time.Duration.between(LocalDateTime.now(), task.getExecuteTime()).toMillis();
        if (delay <= 0) {
            // 如果任务已经到期，立即执行
            executeTask(task);
        } else {
            // 否则安排在到期时间执行
            scheduledExecutorService.schedule(() -> {
                if (task.getStatus() == TaskStatus.PENDING && canExecuteTask(task)) {
                    executeTask(task);
                }
            }, delay, TimeUnit.MILLISECONDS);
        }
    }

    @Override
    public Task getTask(String taskId) {
        return taskStorage.getTask(taskId);
    }

    @Override
    public List<Task> getTasksByStatus(TaskStatus status) {
        return taskStorage.getAllTasks().stream()
                .filter(task -> task.getStatus() == status)
                .collect(Collectors.toList());
    }

    @Override
    public boolean cancelTask(String taskId) {
        Task task = taskStorage.getTask(taskId);
        if (task != null && task.getStatus() == TaskStatus.PENDING) {
            task.setStatus(TaskStatus.CANCELLED);
            taskStorage.updateTask(task);
            return true;
        }
        return false;
    }

    @Override
    public TaskMonitor.TaskStatistics getStatistics() {
        return taskMonitor.getStatistics();
    }

    private boolean canExecuteTask(Task task) {
        // 检查任务依赖是否都已完成
        String[] dependencies = task.getDependencies();
        if (dependencies != null && dependencies.length > 0) {
            return Arrays.stream(dependencies)
                    .allMatch(depId -> {
                        Task depTask = taskStorage.getTask(depId);
                        return depTask != null && depTask.getStatus() == TaskStatus.COMPLETED;
                    });
        }
        return true;
    }

    private void executeTask(Task task) {
        executorService.submit(() -> {
            try {
                task.getMetadata().setStartTime(LocalDateTime.now());
                task.setStatus(TaskStatus.RUNNING);
                taskStorage.updateTask(task);
                taskMonitor.taskStarted(task);

                task.execute();

                task.getMetadata().setEndTime(LocalDateTime.now());
                task.setStatus(TaskStatus.COMPLETED);
                taskStorage.updateTask(task);
                taskMonitor.taskCompleted(task);
            } catch (Exception e) {
                handleTaskFailure(task, e);
            }
        });
    }

    private void handleTaskFailure(Task task, Exception e) {
        task.getMetadata().setEndTime(LocalDateTime.now());
        task.getMetadata().setError(e.getMessage(), getStackTrace(e));
        taskMonitor.taskFailed(task);

        if (task.getCurrentRetries() < task.getMaxRetries()) {
            // 重试任务
            task.incrementRetries();
            task.setStatus(TaskStatus.PENDING);
            task.getMetadata().setStartTime(null);
            task.getMetadata().setEndTime(null);
            taskStorage.updateTask(task);

            // 计算重试时间并重新调度
            LocalDateTime retryTime = LocalDateTime.now().plus(task.getRetryInterval(), java.time.temporal.ChronoUnit.MILLIS);
            taskQueue.addTask(task);
            scheduleTaskExecution(task);
        } else {
            // 达到最大重试次数，标记为失败
            task.setStatus(TaskStatus.FAILED);
            taskStorage.updateTask(task);
        }
    }

    private String getStackTrace(Exception e) {
        StringBuilder sb = new StringBuilder();
        for (StackTraceElement element : e.getStackTrace()) {
            sb.append(element.toString()).append("\n");
        }
        return sb.toString();
    }
} 